from setuptools import setup
import os
import subprocess

# Note on version postfix:
#   Use .devX on the developer's branch
#   Use rcY  on release branch.
# IMPORTANT: The (.) should be present for .dev, but not for rc (PEP 440).
NVBLOX_VERSION_NUMBER = '0.0.8'
NVBLOX_VERSION_PATCH = 'rc5'
NVBLOX_VERSION = f'{NVBLOX_VERSION_NUMBER}{NVBLOX_VERSION_PATCH}'


def get_version_with_platform_tag() -> str:
    """Get version with cuda and ubuntu tags appended"""
    return f'{NVBLOX_VERSION}+cu{get_cuda_major_version()}ubuntu{get_ubuntu_major_version()}'


def get_cuda_major_version() -> str:
    """Get the major CUDA version from the environment variable CUDA_VERSION."""
    return os.environ.get('CUDA_VERSION', '').split('.')[0]


def get_ubuntu_major_version() -> str:
    """Get the major Ubuntu version from the system."""
    return subprocess.check_output(['lsb_release', '-rs']).decode().strip().split('.')[0]


def get_git_sha() -> str:
    """Get the build number of nvblox_torch from git."""
    # To ensure this work in environments where git is not available, we need to obtain
    # this information from the .git dir rather than using the git command
    script_dir = os.path.dirname(os.path.abspath(__file__))
    git_dir = os.path.join(script_dir, '..', '.git')

    # Read the current ref/branch name from .git/HEAD
    try:
        with open(os.path.join(git_dir, 'HEAD'), 'r', encoding='utf-8') as f:
            git_ref = f.read().strip()

            # If HEAD is on a branch, it contains the string 'ref: BRANCH_NAME'
            # We then need to look up the SHA for BRANCH_NAME.
            if git_ref.startswith('ref: '):
                git_ref = git_ref.split(' ')[1]
                # Look up the SHA for the ref/branch name
                with open(os.path.join(git_dir, git_ref), 'r', encoding='utf-8') as f:
                    git_sha = f.read().strip()
            # If there's no branch checked out, the HEAD file contains the current SHA
            else:
                git_sha = git_ref

        assert len(git_sha) == 40, f'git_sha is not 40 characters: {git_sha}'

    # TODO(dtingdahl): aarch64 build fails here as it seems like setup.py is invoked
    # from a temporary directory rather than from the repo.
    except (FileNotFoundError, NotADirectoryError):
        print('Warning: unknown git SHA for nvblox_torch')
        git_sha = 'unknown'

    return git_sha


def get_expected_wheel_filename(build_number: str) -> str:
    """Get the expected wheel filename for the current version/build of nvblox_torch."""
    version = get_version_with_platform_tag()
    return f'nvblox_torch-{version}-{build_number}-py3-none-linux_x86_64.whl'


def create_version_file() -> None:
    """Generate a version file containing build info"""
    version_path = os.path.join('nvblox_torch', '_version.py')

    version_file_content = """# This file is automatically generated by setup.py
__version__ = '{version}'
__git_sha__ = '{git_sha}'
"""
    with open(version_path, 'w', encoding='utf-8') as f:
        f.write(
            version_file_content.format(version=get_version_with_platform_tag(),
                                        git_sha=get_git_sha()))


if __name__ == '__main__':
    # Generate a version file
    create_version_file()
    # Proceed with setup
    setup(version=get_version_with_platform_tag())
